package cn.yuemouren.tree.avl;

import java.util.ArrayList;

/**
 * @Author: Timi
 * @Description: AVL树（平衡二叉树）
 * @Date: 2020/7/13 21:00
 * @Version: v1.0
 */
public class AvlTree<K extends Comparable<K>, V> {

    private class Node {
        public K key;
        public V value;
        public Node left, right;
        public int height;

        public Node(K key) {
            this(key, null);
        }

        public Node(K key, V value) {
            this.key = key;
            this.value = value;
            this.left = null;
            this.right = null;
            this.height = 1;
        }
    }

    private Node root;
    private int size;

    public AvlTree(Node root, int size) {
        this.root = null;
        this.size = 0;
    }

    //获取权重
    public int getHeight(Node node) {
        if (null == node) {
            return 0;
        }
        return node.height;
    }

    //获取平衡因子
    public int getBalanceFactor(Node node) {
        if (null == node) {
            return 0;
        }
        return getHeight(node.left) - getHeight(node.right);
    }

    //判断是否是二分搜索树
    public boolean isBinarySearch() {
        ArrayList<K> arrayList = new ArrayList<>();
        isOrder(arrayList, root);
        for (int i = 1; i < arrayList.size(); i++) {
            if (arrayList.get(i - 1).compareTo(arrayList.get(i)) > 0) {
                return false;
            }
        }
        return true;
    }

    private void isOrder(ArrayList<K> arrayList, Node node) {
        if (null == node) {
            return;
        }
        isOrder(arrayList, node.left);
        arrayList.add(node.key);
        isOrder(arrayList, node.right);
    }

    //判断是否是平衡树
    public boolean isBalanced() {
        return isBalanced(root);
    }

    private boolean isBalanced(Node node) {
        if (null == node) {
            return true;
        }
        int abs = Math.abs(getBalanceFactor(node));
        if (abs > 1) {
            return false;
        }
        return isBalanced(node.left) && isBalanced(node.right);
    }

    /**
     * 右旋转
     * y                            x
     * /  \                        /    \
     * x    T4                     z        y
     * / \                         / \    /    \
     * Z   T3     ------------>    T1 T2   T3    T4
     * / \
     * T1  T2
     *
     * @param y
     *
     * @return
     */
    private Node rightSpin(Node y) {
        Node x = y.left;
        Node T3 = x.right;

        x.right = y;
        y.left = T3;

        y.height = 1 + Math.max(getHeight(y.left), getHeight(y.right));
        x.height = 1 + Math.max(getHeight(x.left), getHeight(x.right));

        return x;
    }

    /**
     * 左旋转
     * 图跟上面的图反过来
     *
     * @param y
     *
     * @return
     */
    private Node leftSpin(Node y) {
        Node x = y.right;
        Node T3 = x.left;

        x.left = y;
        y.right = T3;

        y.height = 1 + Math.max(getHeight(y.left), getHeight(y.right));
        x.height = 1 + Math.max(getHeight(x.left), getHeight(x.right));

        return x;
    }

    /**
     * 添加映射关系
     *
     * @param key
     * @param value
     */

    public void add(K key, V value) {
        root = add(root, key, value);
    }

    private Node add(Node node, K key, V value) {
        if (null == node) {
            size++;
            return node = new Node(key, value);
        }
        if (key.compareTo(node.key) < 0) {
            node.left = add(node.left, key, value);
        } else if (key.compareTo(node.key) > 0) {
            node.right = add(node.right, key, value);
        } else {
            node.value = value;
        }

        node.height = 1 + Math.max(getHeight(node.left), getHeight(node.right));

        int balanceFactor = getBalanceFactor(node);
        if (Math.abs(balanceFactor) > 1) {
            System.out.println("平衡因子为:" + balanceFactor);
        }

        //LL 左子树的左子树
        if (balanceFactor > 1 && getBalanceFactor(node.left) >= 0) {
            return rightSpin(node);
        }

        //RR 右子树的右子树
        if (balanceFactor < -1 && getBalanceFactor(node.right) <= 0) {
            return leftSpin(node);
        }

        //LR 左子树的右子树
        if (balanceFactor > 1 && getBalanceFactor(node.left) < 0) {
            node.left = leftSpin(node.left);
            return rightSpin(node);
        }

        //RL 右子树的左子树
        if (balanceFactor < -1 && getBalanceFactor(node.right) > 0) {
            node.right = rightSpin(node.right);
            return leftSpin(node);
        }

        return node;
    }

    private Node getNode(Node node, K key) {
        if (null == node) {
            return null;
        }
        if (key.compareTo(node.key) < 0) {
            return getNode(node.left, key);
        } else if (key.compareTo(node.key) > 0) {
            return getNode(node.left, key);
        } else {
            return node;
        }
    }

    /**
     * 更改键为key的value
     *
     * @param key
     * @param newValue
     *
     * @return
     */
    public V set(K key, V newValue) {
        Node node = getNode(root, key);
        if (null == node) {
            throw new IllegalArgumentException(key + " not is exist!");
        }
        V oldValue = node.value;
        node.value = newValue;
        return oldValue;
    }

    public boolean contain(K key) {
        return getNode(root, key) != null;
    }

    public Node remove(K key) {
        return remove(root, key);
    }

    private Node remove(Node node, K key) {
        Node resNode = null;
        if (key.compareTo(node.key) < 0) {
            resNode.left = remove(node.left, key);
        } else if (key.compareTo(node.key) > 0) {
            resNode.right = remove(node.right, key);
        } else {//需要删除
            if (null == node.left) {//没有左子树
                Node rightNode = node.right;
                node.right = null;
                size--;
                resNode = rightNode;
            } else if (null == node.right) {//没有右子树
                Node leftNode = node.left;
                node.left = null;
                size--;
                resNode = leftNode;
            } else {//左右子树均存在,两种解决办法，一个是找左子树最大值，还有一个是找右子树最小值
                Node maxNode = getMax(node.left);
                Node removeNode = remove(node, maxNode.key);
                maxNode.left = node.left;
                maxNode.right = node.right;
                resNode = node;
                node.left = node.right = null;
            }
        }

        //维护height和平衡因子
        resNode.height = 1 + Math.max(getHeight(resNode.left), getHeight(resNode.right));

        int balanceFactor = getBalanceFactor(resNode);
        //LL
        if (balanceFactor > 1 && getBalanceFactor(node.left) >= 0) {
            return rightSpin(node);
        }

        //RR
        if (balanceFactor < -1 && getBalanceFactor(node.right) <= 0) {
            return leftSpin(node);
        }

        //LR
        if (balanceFactor > 1 && getBalanceFactor(node.left) < 0) {
            node.left = leftSpin(node.left);
            return rightSpin(node);
        }

        //RL
        if (balanceFactor < -1 && getBalanceFactor(node.right) > 0) {
            node.right = rightSpin(node.right);
            return leftSpin(node);
        }

        return resNode;
    }

    /**
     * 递归实现获取树中最右节点（最大节点）
     *
     * @param node
     *
     * @return
     */
    private Node getMax(Node node) {
        if (null == node.right) {
            return node;
        }
        return getMax(node.right);
    }

    /**
     * 递归实现获取树中最左侧节点（最小节点）
     *
     * @param node
     *
     * @return
     */
    private Node getMin(Node node) {
        if (null == node.left) {
            return node;
        }
        return getMin(node.left);
    }

    public int getSize() {
        return size;
    }
}
